# rsa_core.py

import math_utils # 导入我们自己的数学模块
import random
import hashlib # 用于OAEP中的哈希
import os      # 用于OAEP中的os.urandom生成种子

# --- 密钥生成 (与之前相同) ---
def generate_keypair(bits=2048):
    # ... (此函数内容保持不变，我们只修改加密和解密)
    if bits < 16: 
        raise ValueError("密钥位数太小，至少应为16位")
    p_bits = bits // 2
    q_bits = bits - p_bits
    print(f"正在生成 {p_bits} 位素数 p...")
    p = math_utils.generate_prime_number(p_bits)
    print(f"p = {p}\n")
    print(f"正在生成 {q_bits} 位素数 q...")
    q = math_utils.generate_prime_number(q_bits)
    while p == q: 
        print("p 和 q 相同，重新生成 q...")
        q = math_utils.generate_prime_number(q_bits)
    print(f"q = {q}\n")
    n = p * q
    print(f"模数 n = p * q = {n}")
    print(f"n 的位长: {n.bit_length()}\n")
    phi_n = (p - 1) * (q - 1)
    print(f"欧拉函数 phi_n = {phi_n}\n")
    e = 65537
    # Python 3.9+ math.gcd handles multiple arguments
    # For older versions, or to be explicit:
    current_gcd = math_utils.extended_gcd(e, phi_n)[0] # extended_gcd returns (gcd, x, y)
    while current_gcd != 1:
        e = random.randrange(3, phi_n, 2) 
        print(f"尝试新的 e: {e}")
        current_gcd = math_utils.extended_gcd(e, phi_n)[0]
    print(f"公钥指数 e = {e}\n")
    d = math_utils.mod_inverse(e, phi_n)
    print(f"私钥指数 d = {d}\n")
    public_key = (n, e)
    private_key = (n, d)
    return public_key, private_key


# --- OAEP 辅助常量 ---
# 默认使用 SHA-256
DEFAULT_HASH_FUNC = hashlib.sha256

# --- OAEP 加密 ---
def oaep_encode(message_bytes, k_bytes, label_bytes=b"", hash_func=DEFAULT_HASH_FUNC):
    """
    OAEP 编码 (PKCS#1 v2.2 EME-OAEP-ENCODE)
    message_bytes: 待编码的原始消息字节串
    k_bytes: 模数n的字节长度 (EM的长度)
    label_bytes: 可选的标签字节串
    hash_func: 使用的哈希算法 (例如 hashlib.sha256)
    """
    m_len = len(message_bytes)
    h_len = hash_func().digest_size

    # 步骤 1a: 标签L过长检查 (实际不太可能触发)
    # if len(label_bytes) > hash_func().block_size: # 这是一个宽松的上限
    #     raise ValueError("Label is too long")

    # 步骤 1b: 消息长度检查
    # EM = 0x00 || maskedSeed || maskedDB
    # maskedDB = DB XOR MGF1(seed, k - hLen - 1)
    # DB = lHash || PS || 0x01 || M
    # len(DB) = hLen + len(PS) + 1 + mLen
    # len(maskedSeed) = hLen
    # len(EM) = 1 + hLen + len(maskedDB) = k_bytes
    # len(maskedDB) = k_bytes - hLen - 1
    # 所以 len(DB) = k_bytes - hLen - 1
    # k_bytes - hLen - 1 = hLen + len(PS) + 1 + mLen
    # len(PS) = k_bytes - 2*hLen - 2 - mLen
    # 要求 len(PS) >= 0
    if m_len > k_bytes - 2 * h_len - 2:
        raise ValueError(f"Message too long for OAEP padding. Max length: {k_bytes - 2 * h_len - 2}, got: {m_len}")

    # 步骤 2: lHash = H(L)
    l_hash = hash_func(label_bytes).digest()

    # 步骤 3: 生成 PS，一个包含 k - mLen - 2hLen - 2 个零字节的字节串
    ps_len = k_bytes - m_len - 2 * h_len - 2
    ps = b'\x00' * ps_len

    # 步骤 4: DB = lHash || PS || 0x01 || M
    db = l_hash + ps + b'\x01' + message_bytes
    # print(f"DB length: {len(db)}, expected: {k_bytes - h_len - 1}")
    assert len(db) == k_bytes - h_len - 1

    # 步骤 5: 生成随机种子 seed (长度为 hLen)
    seed = os.urandom(h_len)

    # 步骤 6: dbMask = MGF(seed, k - hLen - 1)
    db_mask = math_utils.mgf1(seed, k_bytes - h_len - 1, hash_func)

    # 步骤 7: maskedDB = DB XOR dbMask
    masked_db = math_utils.bytes_xor(db, db_mask)

    # 步骤 8: seedMask = MGF(maskedDB, hLen)
    seed_mask = math_utils.mgf1(masked_db, h_len, hash_func)

    # 步骤 9: maskedSeed = seed XOR seedMask
    masked_seed = math_utils.bytes_xor(seed, seed_mask)

    # 步骤 10: EM = 0x00 || maskedSeed || maskedDB
    em = b'\x00' + masked_seed + masked_db
    # print(f"EM length: {len(em)}, expected k_bytes: {k_bytes}")
    assert len(em) == k_bytes
    
    return em

def encrypt_oaep(public_key, plaintext_bytes, label_bytes=b"", hash_func=DEFAULT_HASH_FUNC):
    """
    使用公钥和OAEP填充加密字节串消息
    public_key: (n, e)
    plaintext_bytes: 待加密的字节串消息
    label_bytes: OAEP的可选标签
    hash_func: OAEP使用的哈希函数
    """
    n, e = public_key
    k_bytes = (n.bit_length() + 7) // 8 # 模数n的字节长度

    # OAEP编码
    encoded_message = oaep_encode(plaintext_bytes, k_bytes, label_bytes, hash_func)
    
    # 将编码后的消息字节串转换为整数 m
    m_encoded_int = math_utils.bytes_to_int(encoded_message)

    # 密文 c = m_encoded_int^e mod n
    c = math_utils.power(m_encoded_int, e, n)
    return c # 返回整数形式的密文


# --- OAEP 解密 ---
def oaep_decode(encoded_message_bytes, k_bytes, label_bytes=b"", hash_func=DEFAULT_HASH_FUNC):
    """
    OAEP 解码 (PKCS#1 v2.2 EME-OAEP-DECODE)
    encoded_message_bytes (EM): 从RSA解密得到的字节串 (长度应为k_bytes)
    k_bytes: 模数n的字节长度
    label_bytes: 可选的标签字节串
    hash_func: 使用的哈希算法
    返回: 原始消息字节串，如果解码失败则抛出 ValueError
    """
    em = encoded_message_bytes
    h_len = hash_func().digest_size

    # 步骤 1a: 标签L过长检查 (同编码)
    # if len(label_bytes) > hash_func().block_size:
    #     raise ValueError("Decryption error (label too long check failed - unlikely)")

    # 步骤 1b: EM长度检查 (应为k_bytes)
    if len(em) != k_bytes:
        raise ValueError("Decryption error (EM length incorrect)")
    
    # 步骤 1c: EM太短检查
    # EM = Y || maskedSeed || maskedDB. Y是第一个字节 (0x00)
    # maskedSeed 长度 hLen, maskedDB 长度 k - hLen - 1
    # 所以 EM 最小长度 1 + hLen + (hLen + 1) (当 PS 为空, M 为空)
    # k >= 2*hLen + 2
    if k_bytes < 2 * h_len + 2:
        raise ValueError("Decryption error (k too small for hash function)")

    # 步骤 2: lHash = H(L)
    l_hash_expected = hash_func(label_bytes).digest()
    
    # 步骤 3: 分离 EM
    y = em[0:1]          # 第一个字节
    masked_seed = em[1 : 1 + h_len]
    masked_db = em[1 + h_len :]

    if len(masked_seed) != h_len or len(masked_db) != k_bytes - h_len - 1:
        # 这一步理论上不会出错如果 len(em) == k_bytes 且 k_bytes >= 1 + h_len
        raise ValueError("Decryption error (EM component lengths invalid)")

    # 步骤 4: seedMask = MGF(maskedDB, hLen)
    seed_mask = math_utils.mgf1(masked_db, h_len, hash_func)

    # 步骤 5: seed = maskedSeed XOR seedMask
    seed = math_utils.bytes_xor(masked_seed, seed_mask)

    # 步骤 6: dbMask = MGF(seed, k - hLen - 1)
    db_mask = math_utils.mgf1(seed, k_bytes - h_len - 1, hash_func)

    # 步骤 7: DB = maskedDB XOR dbMask
    db = math_utils.bytes_xor(masked_db, db_mask)
    # print(f"DB length: {len(db)}, expected: {k_bytes - h_len - 1}")
    assert len(db) == k_bytes - h_len - 1

    # 步骤 8: 分离 DB = lHash' || PS || 0x01 || M'
    l_hash_prime = db[0:h_len]

    # 验证 lHash'
    if l_hash_prime != l_hash_expected:
        raise ValueError("Decryption error (lHash mismatch)")

    # 寻找 0x01 分隔符，并验证 PS
    # PS 从 lHash_prime 之后开始，直到 0x01
    separator_index = -1
    for i in range(h_len, len(db)):
        if db[i] == 0x01:
            separator_index = i
            break
        if db[i] != 0x00: # PS 中间出现非零字节
            raise ValueError("Decryption error (non-zero byte in PS)")
    
    if separator_index == -1: # 没有找到 0x01 分隔符
        raise ValueError("Decryption error (0x01 separator not found)")

    # 验证第一个字节 Y
    if y != b'\x00':
        raise ValueError("Decryption error (leading byte Y is not 0x00)")
        
    # 原始消息 M'
    message_bytes = db[separator_index + 1 :]
    return message_bytes


def decrypt_oaep(private_key, ciphertext_int, label_bytes=b"", hash_func=DEFAULT_HASH_FUNC):
    """
    使用私钥和OAEP填充解密整数形式的密文
    private_key: (n, d)
    ciphertext_int: 整数形式的密文
    label_bytes: OAEP的可选标签
    hash_func: OAEP使用的哈希函数
    """
    n, d = private_key
    k_bytes = (n.bit_length() + 7) // 8 # 模数n的字节长度

    # 检查密文是否在 [0, n-1] 范围内
    if not (0 <= ciphertext_int < n):
        raise ValueError("Decryption error (ciphertext out of range)")

    # 明文 m_encoded_int = ciphertext_int^d mod n
    m_encoded_int = math_utils.power(ciphertext_int, d, n)
    
    # 将解密后的整数转换回字节串 EM，长度必须是 k_bytes
    # math_utils.int_to_bytes 需要能处理前导零字节的情况
    encoded_message_bytes = math_utils.int_to_bytes(m_encoded_int, min_length=k_bytes)
    
    # OAEP解码
    try:
        plaintext_bytes = oaep_decode(encoded_message_bytes, k_bytes, label_bytes, hash_func)
    except ValueError as e:
        # 重要: 对于OAEP，解码失败时不能泄露具体原因，通常返回一个通用的解密错误
        # 在这里，我们为了调试会打印具体错误，但实际应用中应捕获并抛出通用错误
        # print(f"OAEP Decode Error: {e}")
        raise ValueError("Decryption error (OAEP decoding failed)") # 抛出通用错误

    return plaintext_bytes

